#include <fstream>
#include <random>

#include <bthread/bthread.h>
#include <gtest/gtest.h>

#include "util/fileaio/aio_file_reader.h"

DEFINE_string(file_read_testfile, "", "");
DEFINE_uint32(file_test_read_ioring_queue_size, 256, "");
DEFINE_uint32(file_test_read_concurrency, 16, "");
DEFINE_int32(file_read_test_tasks, 128, "num of threads or bthreads");

using namespace hawking::indexlib;
const size_t test_file_size = 1024 * 1024 * 1024;
std::uniform_int_distribution<uint32_t> distr(0, test_file_size - 256);
std::default_random_engine eng{std::random_device{}()};

void MakeTestFile() {
    struct stat st;
    if (::stat(FLAGS_file_read_testfile.data(), &st) == 0) {
        return;
    }
    std::ofstream ofs(FLAGS_file_read_testfile.data());
    for (size_t idx = 0; idx < test_file_size; ++idx) {
        ofs << static_cast<char>(idx);
    }
    ofs.close();
}

void ThreadTestFunction(DiskFileReadProcesser* processor) {
    uint8_t buffer[256];
    while (1) {
        size_t offset = static_cast<size_t>(distr(eng));
        auto ret = processor->Read(buffer, 256, offset);
        EXPECT_EQ(int(Status::OK), int(ret.Code()));
    }
}

void* BthreadTestFunction(void* args) {
    DiskFileReadProcesser* processor = reinterpret_cast<DiskFileReadProcesser*>(args);
    uint8_t* buffer = new uint8_t[256];
    while (1) {
        size_t offset = static_cast<size_t>(distr(eng));
        auto ret = processor->Read(buffer, 256, offset);
        EXPECT_EQ(int(Status::OK), int(ret.Code()));
    }

    delete[] buffer;
    return nullptr;
}

AioOption MakeAioOption() {
    AioOption opt;
    return opt;
}

TEST(file_read_test, multi_thread_test) {
    MakeTestFile();
    DiskFileReadProcesser processor;
    AioOption opt;
    bool init_ret = processor.Init(opt, FLAGS_file_read_testfile.data(), true);
    EXPECT_EQ(true, init_ret);
    if (!init_ret) {
        return;
    }
    
    std::vector<std::thread> reader_threads;
    reader_threads.reserve(FLAGS_file_read_test_tasks);
    for (int32_t _ = 0; _ < FLAGS_file_read_test_tasks; ++_) {
        reader_threads.emplace_back(ThreadTestFunction, &processor);
    }

    for (auto& t : reader_threads) {
        t.join();
    }
}

TEST(file_read_test, multi_bthread_test) {
    MakeTestFile();
    DiskFileReadProcesser processor;
    AioOption opt;
    bool init_ret = processor.Init(opt, FLAGS_file_read_testfile.data(), true);
    EXPECT_EQ(true, init_ret);
    if (!init_ret) {
        return;
    }

    std::vector<bthread_t> bids;
    bids.reserve(FLAGS_file_read_test_tasks);
    for (int32_t _ = 0; _ < FLAGS_file_read_test_tasks; ++_) {
        bids.emplace_back();
        bthread_t& bid = bids.back();
        bthread_start_background(&bid, nullptr, BthreadTestFunction, &processor);
    }

    for (auto& bid : bids) {
        bthread_join(bid, nullptr);
    }
}

int main(int argc, char* argv[]) {
    ::testing::InitGoogleTest(&argc, argv);
    ::google::ParseCommandLineFlags(&argc, &argv, true);
    srand(static_cast<unsigned>(time(0)));
    return RUN_ALL_TESTS();
}
